#include <bezier_predict/predictor.h>
#include <stdio.h>
#include <ros/ros.h>
#include <ros/console.h>
#include <iostream>
#include <fstream>
#include <string>
#include <math.h>
#include <geometry_msgs/PointStamped.h>

using namespace std;
using namespace Eigen;
using namespace ego_planner;

namespace ego_planner {

void Predictor::init(ros::NodeHandle &nh,
                     Visualization::Ptr vis,
                     ObsPosesPtr pos_ptr,
                     ObsTrajsPtr traj_ptr)
{
  visualization_ = vis;
  obs_poses_ = pos_ptr;
  obs_trajs_ = traj_ptr;

  nh.param("prediction/time_remain_", time_remain_, -1.0);
  nh.param("prediction/time_vanish_", time_vanish_, -1.0);
  nh.param("prediction/time_turncate_", time_turncate_, 4.0);
  nh.param("prediction/time_scale_", time_scale_, -1.0);

  vis_predict_list_pub_ =
      nh.advertise<visualization_msgs::Marker>("predict/points_list", 2);
  vis_control_pt_pub_ =
      nh.advertise<visualization_msgs::Marker>("predict/control_pts", 2);
  manager_timer_ = nh.createTimer(ros::Duration(0.2), &Predictor::managerCb, this);

  // set param for bezier predictor
  bezierPredict_.setSegInterval((int)list_min_size_, (int)list_max_size_);
  bezierPredict_.setSampleInterval(sample_interval_);
  bezierPredict_.setTimeScale(time_scale_);
}

void Predictor::managerCb(const ros::TimerEvent &e)
{
  ros::Time t_now = ros::Time::now();

  have_obs_ = false;

  ObjsMap::iterator it;
  for (it = predicted_objs_.begin(); it != predicted_objs_.end(); it++)
  {
    if (!it->second.target_list_.empty()) {
      double t = t_now.toSec() - it->second.target_list_.back()(3);
      if (t > time_vanish_)
      {
        it->second.target_list_.clear();
        it->second.have_pos_ = false;
      }
    }

    have_obs_ = have_obs_ || it->second.have_pos_;

    // turncate the time is too older
    for (size_t i = 0; i < it->second.target_list_.size(); i++)
    {
      if (t_now.toSec() - it->second.target_list_[i](3) < time_turncate_)
      {
        if (i == 0) break;
        it->second.target_list_.erase(it->second.target_list_.begin(),
                                      it->second.target_list_.begin() + i);
        break;
      }
    }
  }

  // clear timeout trajs
  for (size_t i = 0; i < obs_trajs_->size(); i++)
  {
    double t = (t_now - obs_trajs_->at(i)->start_time_).toSec();
    if (t > obs_trajs_->at(i)->duration_ + time_vanish_)
    {
      clearObsTraj(i);
      i--;
    }
  }
}

void Predictor::setInput(const std::vector<int> &_ids,
                         const std::vector<Eigen::Vector4d> &_targets)
{
  if (_ids.empty()) return;

  have_obs_ = true;

  for (size_t i = 0; i < _ids.size(); i++)
  {
    ObjsMap::iterator it = predicted_objs_.find(_ids[i]);
    if (it != predicted_objs_.end())
    {
      it->second.setInput(_targets[i]);
    }
    else {
      predicted_objs_[_ids[i]] = PredictedObj(list_max_size_);
      predicted_objs_[_ids[i]].setInput(_targets[i]);
    }

    updateObsPoses(_ids[i], _targets[i]);
  }

  updatePredict();
}

/* update prediction for each obj with enough data */
void Predictor::updatePredict()
{
  std::vector<Eigen::Vector3d> samples;

  ObjsMap::iterator it;
  for (it = predicted_objs_.begin(); it != predicted_objs_.end(); it++)
  {
    if (it->second.target_list_.size() < list_min_size_)
      continue;

    bool isSuccess = bezierPredict_.TrackingGeneration(
        5, 5, it->second.target_list_);

    if (!isSuccess) continue;

    it->second.t_start_gl_ = bezierPredict_.t_predict_start_gl_;
    it->second.R2_ = bezierPredict_.R2_;
    it->second.prop_start_ = bezierPredict_.prop_start_;
    it->second.Evaluation_ = bezierPredict_.Evaluation_;

    std::vector<Eigen::Vector3d>
        sample = bezierPredict_.getPredictPosList();

    samples.insert(samples.end(), sample.begin(), sample.end());

    Eigen::MatrixXd ctrl_pts;
    std::vector<Eigen::Vector3d> start_end_derivatives;
    start_end_derivatives.push_back(bezierPredict_.getStartVel());
    start_end_derivatives.push_back(Eigen::Vector3d::Zero());
    start_end_derivatives.push_back(bezierPredict_.getStartAcc());
    start_end_derivatives.push_back(Eigen::Vector3d::Zero());
    UniformBspline::BsplineFitting(sample_interval_, sample,
                                   start_end_derivatives, ctrl_pts);

    ObsTraj obs_traj;
    obs_traj.id = it->first;
    obs_traj.traj_pos_ = UniformBspline(ctrl_pts, 3, sample_interval_);
    obs_traj.start_time_ = ros::Time(bezierPredict_.t_predict_start_gl_);
    obs_traj.duration_ = obs_traj.traj_pos_.getTimeSum();

    updateObsTrajs(it, obs_traj); // record the traj
  }

  Eigen::Vector4d color(0, 1, 0.5, 1);
  visualization_->displaySphereList(vis_predict_list_pub_, samples,
                                    0.2, color, 10);
}

/**
 * @param pos
 * @param R2 : goodness of fit, == -5.0 if no traj
 * @param times
 * @param is_dangerous: if have no traj, dangerous, else, get from traj
 */
bool Predictor::getPredictPos(std::vector<std::vector<Eigen::Vector3d>> &pos,
                              std::vector<std::vector<double>> &R2,
                              std::vector<std::vector<double>> &times,
                              std::vector<bool> &is_dangerous)
{
  double t_now_gl = ros::Time::now().toSec();

  pos.clear();
  R2.clear();
  times.clear();
  is_dangerous.clear();

  ObjsMap::iterator it;
  for (it = predicted_objs_.begin(); it != predicted_objs_.end(); it++)
  {
    if (it->second.have_traj_)
    {
      double t_now = t_now_gl - it->second.traj_->start_time_.toSec();
      double t_end = it->second.traj_->duration_ * 2.0 / 3.0 + time_remain_;
      if (t_now < t_end)
      {
        is_dangerous.push_back(it->second.traj_->is_dangerous_);

        std::vector<Eigen::Vector3d> _pos;
        std::vector<double> _R2;
        std::vector<double> _times;

        for (double t = 0.0; t + t_now < t_end; t += 1.0)
        {
          _pos.push_back(it->second.traj_->traj_pos_.getPos(t + t_now));
          _R2.push_back(it->second.getUncertainty(t + t_now_gl));
          _times.push_back(t + t_now_gl);
        }
        pos.push_back(_pos);
        R2.push_back(_R2);
        times.push_back(_times);

        continue;
      }
      else {
        it->second.have_traj_ = false;
      }
    }

    if (it->second.have_pos_)
    {
      double t_now = t_now_gl - it->second.target_list_.back()(3);
      if (t_now < time_remain_) {
        is_dangerous.push_back(true); // if have no traj, dangerous
        pos.push_back({it->second.target_list_.back().head(3)});
        R2.push_back({-5.0}); // goodness of fit, == -5.0 if no traj
        times.push_back({t_now_gl});
        continue;
      }
      else {
        it->second.have_pos_ = false;
      }
    }
  }

  if (pos.empty()) return false;

  return true;
}

}